'''
Author: 渔火Arcadia  https://github.com/yhArcadia
Date: 2025-04-18 00:13:00
LastEditors: 渔火Arcadia
LastEditTime: 2025-04-18 22:42:55
FilePath: /example/model.py
Description: 

Copyright (c) 2025 by 渔火Arcadia 1761869682@qq.com, All Rights Reserved. 
'''
import torch
import torch.nn as nn

class mymodel(nn.Module):
    def __init__(self):
        super().__init__()
        self.卷积= nn.Conv2d(3, 64, kernel_size=3, padding=1) #定义卷积层，输入通道数为3，输出通道数为64，卷积核大小为3，填充为1
        self.展平 = nn.Flatten(start_dim=1) #定义展平层，将卷积层输出的特征图展平成通道数*长*宽维的向量
        self.映射到4维向量=nn.Linear(64*6*6, 4) #定义全连接层，输入特征数为64*6*6，输出特征数为4
        

    def forward(self, x): #input: [batchsize, 3, 6, 6]
        
        # 卷积层，输入为3通道（RGB）的图片，输出为64通道的特征图
        x = self.卷积(x) #output: [batchsize, 64, 6, 6]
        
        # 将特征图展平成：通道数*长*宽=64*6*6=2304维 的向量
        x = self.展平(x) #output: [batchsize, 2304]
        
        # 将2304维的特征向量映射到4维，因为数据集中一共有4个类别
        x = self.映射到4维向量(x) #output: [batchsize, 4]
        
        # 对4维向量进行softmax操作，得到4个分类的概率分布
        x = torch.softmax(x,dim=1)#output: [batchsize, 4] 
        
        return x